/******************************************************
 * @file mnist.c
 * @author Destiny 
 * @brief mnist模型定义
 * 
 * @version 0.1
 * @date 2024-07-18
 *****************************************************/
/* Includes ------------------------------------------------------------------*/
#include <stdio.h>
#include <stdlib.h>
/* Private Includes ----------------------------------------------------------*/
#include "config.h"
#include "model_manager.h"

static PT_InoutputData Mnist(PT_InoutputData ptImageData);
static int MnistParamsInit(void);
static void MnistParamsFree(void);

/*神经网络卷积核参数*/
static T_KernelParams g_tMnistKernelParams[] = {
    {
        .name           = "Conv2d_Kernel1",
        .iNum           = 16,                  /*数量*/
        .iDim           = 1,			        /*维数*/
        .iRow           = 5,                   /*行*/
        .iCol           = 5, 
        .iBiasLength    = 16,
        .ParamsPathInH5 = "/conv2d_3/conv2d_3/"
    },
    {
        .name           = "Conv2d_Kernel2",
        .iNum           = 8,                  /*数量*/
        .iDim           = 16,			        /*维数*/
        .iRow           = 5,                   /*行*/
        .iCol           = 5, 
        .iBiasLength    = 8,
        .ParamsPathInH5 = "/conv2d_4/conv2d_4/"
    },
    {
        .name           = "Dense1",
        .iNum           = 1,                  /*数量*/
        .iDim           = 1,			        /*维数*/
        .iRow           = 392,                   /*行*/
        .iCol           = 200, 
        .iBiasLength    = 200,
        .ParamsPathInH5 = "/dense_3/dense_3/"
    },
    {
        .name           = "Dense2",
        .iNum           = 1,                  /*数量*/
        .iDim           = 1,			        /*维数*/
        .iRow           = 200,                   /*行*/
        .iCol           = 10, 
        .iBiasLength    = 10,
        .ParamsPathInH5 = "/dense_4/dense_4/"
    },
};

/*模型描述结构体*/
static T_ModelDisc g_tMnistModel = {
    .name               = "mnist",
    .pcModelPath        = MODEL_PATH,
    .iNumOfLayer        = 4,
    .ptKernelParamsHead = g_tMnistKernelParams,
    .ModelParamsInit    = MnistParamsInit,
    .ModelParamsFree    = MnistParamsFree,
    .ModelFunction      = Mnist,
};

/*Model: "mnist"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_3 (Conv2D)            (None, 28, 28, 16)        416       
_________________________________________________________________
max_pooling2d_3 (MaxPooling2 (None, 14, 14, 16)        0         
_________________________________________________________________
conv2d_4 (Conv2D)            (None, 14, 14, 8)         3208      
_________________________________________________________________
max_pooling2d_4 (MaxPooling2 (None, 7, 7, 8)           0         
_________________________________________________________________
flatten_2 (Flatten)          (None, 392)               0         
_________________________________________________________________
dense_3 (Dense)              (None, 200)               78600     
_________________________________________________________________
dropout_2 (Dropout)          (None, 200)               0         
_________________________________________________________________
dense_4 (Dense)              (None, 10)                2010      
=================================================================
Total params: 84,234
Trainable params: 84,234
Non-trainable params: 0
_________________________________________________________________*/

/**
 * @brief Mnist网络结构
 * 
 * @param ptImageData 输入的原始图像数据
 * @return PT_InoutputData 输出识别结果
 */
static PT_InoutputData Mnist(PT_InoutputData ptImageData)
{
    int iPadding = (g_tMnistKernelParams[0].iRow - 1)/2;
    /*第一层卷积层*/
    PT_InoutputData ptOutputData1  = Convolution2D(ptImageData,&g_tMnistKernelParams[0],1,iPadding);
    ActivationRelu(ptOutputData1);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData1:%d %d %d\r\n",ptOutputData1->iDim,ptOutputData1->iRow,ptOutputData1->iCol);

    /*第一层池化层*/
    PT_InoutputData ptOutputData2 = PoolOperation(ptOutputData1,2,0,2,1);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData2:%d %d %d\r\n",ptOutputData2->iDim,ptOutputData2->iRow,ptOutputData2->iCol);

    /*第二层卷积层*/
    iPadding = (g_tMnistKernelParams[1].iRow - 1)/2;
    PT_InoutputData ptOutputData3 = Convolution2D(ptOutputData2,&g_tMnistKernelParams[1],1,iPadding);
    ActivationRelu(ptOutputData3);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData3:%d %d %d\r\n",ptOutputData3->iDim,ptOutputData3->iRow,ptOutputData3->iCol);

    /*第二层池化层*/
    PT_InoutputData ptOutputData4 = PoolOperation(ptOutputData3,2,0,2,1);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData4:%d %d %d\r\n",ptOutputData4->iDim,ptOutputData4->iRow,ptOutputData4->iCol);

    /*flatten*/
    PT_InoutputData ptOutputData5 = Flatten(ptOutputData4);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData5:%d %d %d\r\n",ptOutputData5->iDim,ptOutputData5->iRow,ptOutputData5->iCol);

    /*第一层全连接层*/
    PT_InoutputData ptOutputData6 = Dense(ptOutputData5,&g_tMnistKernelParams[2]);
    ActivationTanh(ptOutputData6);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData6:%d %d %d\r\n",ptOutputData6->iDim,ptOutputData6->iRow,ptOutputData6->iCol);

    /*第二层全连接层*/
    PT_InoutputData ptOutputData7 = Dense(ptOutputData6,&g_tMnistKernelParams[3]);
    ActivationSoftmax(ptOutputData7);
    DBG_PRINTF(DLOG_LVL_DEBUG,DLOG_TAG,"ptOutputData7:%d %d %d\r\n",ptOutputData7->iDim,ptOutputData7->iRow,ptOutputData7->iCol);
    return ptOutputData7;
}

/**
 * @brief Mnist参数初始化
 * 
 * @return int 
 */
static int MnistParamsInit(void)
{
    return ModelParamsInit(&g_tMnistKernelParams[0],g_tMnistModel.iNumOfLayer,g_tMnistModel.pcModelPath);
}

/**
 * @brief 释放模型参数内存
 * 
 */
static void MnistParamsFree(void)
{
    ModelParamsFree(&g_tMnistKernelParams[0],g_tMnistModel.iNumOfLayer);
}

/**
 * @brief 注册
 */
int MnisInit(void)
{
    return RegisterModelDisc(&g_tMnistModel);
}